package org.example.cases.single;

import com.antgroup.geaflow.api.function.base.FilterFunction;
import com.antgroup.geaflow.api.function.base.MapFunction;
import com.antgroup.geaflow.api.graph.PGraphWindow;
import com.antgroup.geaflow.api.graph.compute.VertexCentricCompute;
import com.antgroup.geaflow.api.graph.function.vc.VertexCentricCombineFunction;
import com.antgroup.geaflow.api.graph.function.vc.VertexCentricComputeFunction;
import com.antgroup.geaflow.api.pdata.stream.window.PWindowSource;
import com.antgroup.geaflow.api.window.impl.AllWindow;
import com.antgroup.geaflow.common.config.Configuration;
import com.antgroup.geaflow.env.Environment;
import com.antgroup.geaflow.env.EnvironmentFactory;
import com.antgroup.geaflow.model.graph.edge.IEdge;
import com.antgroup.geaflow.model.graph.edge.impl.ValueEdge;
import com.antgroup.geaflow.model.graph.vertex.IVertex;
import com.antgroup.geaflow.model.graph.vertex.impl.ValueVertex;
import com.antgroup.geaflow.pipeline.IPipelineResult;
import com.antgroup.geaflow.pipeline.Pipeline;
import com.antgroup.geaflow.pipeline.PipelineFactory;
import com.antgroup.geaflow.pipeline.task.PipelineTask;
import com.antgroup.geaflow.view.GraphViewBuilder;
import com.antgroup.geaflow.view.IViewDesc.BackendType;
import com.antgroup.geaflow.view.graph.GraphViewDesc;
import org.example.config.JobConfig;
import org.example.function.AbstractVcFunc;
import org.example.function.CSVFileSink;
import org.example.function.CSVFileSource;
import org.example.util.FileUtils;

import java.util.Arrays;
import java.util.Collections;
import java.util.HashMap;
import java.util.Iterator;
import java.util.List;
import java.util.Map;

public class Case2Single {

    private final String snapshotPath;

    private final String resultPath;

    public Case2Single(String snapshotPath, String resultPath) {
        this.snapshotPath = snapshotPath;
        this.resultPath = resultPath;
    }

    public static void main(String[] args) {
        long st = System.currentTimeMillis();
        new Case2Single(args[0], args[1]).start();
        long ed = System.currentTimeMillis();
        System.out.println("cost time: " + (ed - st) + " ms");
    }

    public void start() {
        Environment environment = EnvironmentFactory.onLocalEnvironment();
        IPipelineResult<?> result = submit(environment);
        result.get();
        environment.shutdown();
    }

    public IPipelineResult<?> submit(Environment environment) {
        Pipeline pipeline = PipelineFactory.buildPipeline(environment);
        Configuration envConfig = environment.getEnvironmentContext().getConfig();
        envConfig.put(CSVFileSink.OUTPUT_DIR, this.resultPath);

        pipeline.submit((PipelineTask) pipelineTaskCxt -> {
            Configuration conf = pipelineTaskCxt.getConfig();

            CSVFileSource<IVertex<Long, Integer>> vSource = new CSVFileSource<>(
                    FileUtils.pathJoin(this.snapshotPath, "Account.csv"),
                    Collections.singletonList("accountId"),
                    (CSVFileSource.CSVLineParser<IVertex<Long, Integer>>) line -> {
                        IVertex<Long, Integer> vertex = new ValueVertex<>(Long.valueOf(line.get(0)), 0);
                        return Collections.singletonList(vertex);
                    }
            );
            PWindowSource<IVertex<Long, Integer>> vertices = pipelineTaskCxt
                    .buildSource(vSource, AllWindow.getInstance())
                    .withParallelism(conf.getInteger(JobConfig.SOURCE_PARALLELISM));

            CSVFileSource<IEdge<Long, Integer>> eSource = new CSVFileSource<>(
                    FileUtils.pathJoin(this.snapshotPath, "AccountTransferAccount.csv"),
                    Arrays.asList("fromId", "toId"),
                    (CSVFileSource.CSVLineParser<IEdge<Long, Integer>>) line -> {
                        IEdge<Long, Integer> edge = new ValueEdge<>(
                                Long.valueOf(line.get(0)), Long.valueOf(line.get(1)), 0);
                        return Collections.singletonList(edge);
                    }
            );
            PWindowSource<IEdge<Long, Integer>> edges = pipelineTaskCxt.buildSource(eSource, AllWindow.getInstance())
                    .withParallelism(conf.getInteger(JobConfig.SINK_PARALLELISM));

            int iterationParallelism = conf.getInteger(JobConfig.ITERATOR_PARALLELISM);
            GraphViewDesc graphViewDesc = GraphViewBuilder
                    .createGraphView(GraphViewBuilder.DEFAULT_GRAPH)
                    .withShardNum(iterationParallelism)
                    .withBackend(BackendType.Memory)
                    .build();

            PGraphWindow<Long, Integer, Integer> graphWindow =
                    pipelineTaskCxt.buildWindowStreamGraph(vertices, edges, graphViewDesc);

            CSVFileSink<String[]> csvFileSink = new CSVFileSink<>("result2.csv");
            graphWindow.compute(new PRAlgorithms(3))
                    .compute(iterationParallelism)
                    .getVertices()
                    .filter((FilterFunction<IVertex<Long, Integer>>) record -> record.getValue() > 0)
                    .map((MapFunction<IVertex<Long, Integer>, String[]>) value -> {
                        Long id = value.getId();
                        String vs = String.format("%d", value.getValue());
                        return new String[]{String.valueOf(id), vs};
                    })
                    .sink(csvFileSink)
                    .withParallelism(conf.getInteger(JobConfig.SINK_PARALLELISM));
        });
        return pipeline.execute();
    }

    public static class PRAlgorithms extends VertexCentricCompute<Long, Integer, Integer, Long> {

        public PRAlgorithms(long iterations) {
            super(iterations);
        }

        @Override
        public VertexCentricComputeFunction<Long, Integer, Integer, Long> getComputeFunction() {
            return new PRVertexCentricComputeFunction();
        }

        @Override
        public VertexCentricCombineFunction<Long> getCombineFunction() {
            return null;
        }
    }

    public static class PRVertexCentricComputeFunction extends AbstractVcFunc<Long, Integer, Integer, Long> {

        @Override
        public void compute(Long vertexId, Iterator<Long> messageIterator) {
            long iterationId = this.context.getIterationId();

            if (iterationId == 1) {
                // set init count = 0
                this.context.setNewVertexValue(0);
                List<IEdge<Long, Integer>> outEdges = this.context.edges().getOutEdges();
                for (IEdge<Long, Integer> outEdge : outEdges) {
                    Long targetId = outEdge.getTargetId();
                    this.context.sendMessage(targetId, vertexId);
                }
            } else if (iterationId == 2) {
                while (messageIterator.hasNext()) {
                    Long msg = messageIterator.next();
                    List<IEdge<Long, Integer>> outEdges = this.context.edges().getOutEdges();
                    for (IEdge<Long, Integer> outEdge : outEdges) {
                        Long targetId = outEdge.getTargetId();
                        this.context.sendMessage(targetId, msg);
                    }
                }
            } else if (iterationId == 3) {
                Map<Long, Integer> map = new HashMap<>();
                List<IEdge<Long, Integer>> outEdges = this.context.edges().getOutEdges();
                for (IEdge<Long, Integer> outEdge : outEdges) {
                    Long targetId = outEdge.getTargetId();
                    map.put(targetId, map.getOrDefault(targetId, 0) + 1);
                }
                int cnt = 0;
                while (messageIterator.hasNext()) {
                    Long msg = messageIterator.next();
                    if (map.containsKey(msg)) {
                        cnt += map.get(msg);
                    }
                }
                this.context.setNewVertexValue(cnt);
            }
        }
    }
}
